import zarr
import argparse
import sys
from numprint import p
parser = argparse.ArgumentParser()
parser.add_argument("zarr_path", type=str)
parser.add_argument("ts", type=int,nargs='+')
args = parser.parse_args()

data = zarr.open(args.zarr_path, mode="r")
for ts in args.ts:
    p.print(f"episode_ends[{ts}]",data['meta']['episode_ends'][ts])
    p.print(f"eelocal_state[{ts}]",data['eelocal']['state'][ts])
    p.print(f"qpos_state[{ts}]",data['qpos']['state'][ts])
    p.print(f"qpos_action[{ts}]",data['qpos']['action'][ts])